# -*- coding: utf-8 -*-
# @Time:    2023-05-12 14:35
# @Author:  geng
# @Email:   yonglonggeng@163.com
# @WeChat:  superior_god
# @File:    01screen_pic.py
# @Project: bbox_detection_ai
# @Package: 
# @Ref:
import shutil

import cv2
import os


def video_to_frames(video_path, outPutDirName, img_pre):
    times = 0

    # 提取视频的频率，每10帧提取一个
    frame_frequency = 1

    # 如果文件目录不存在则创建目录
    if not os.path.exists(outPutDirName):
        os.makedirs(outPutDirName)

    # 读取视频帧
    camera = cv2.VideoCapture(video_path)

    while True:
        times = times + 1
        res, image = camera.read()
        if not res:
            print('not res , not image')
            break
        # 按照设置间隔存储视频帧
        if times % frame_frequency == 0:
            cv2.imwrite(outPutDirName + '\\' + img_pre + "_" + str(times) + '.jpg', image)

    print('图片提取结束')
    # 释放摄像头设备
    camera.release()


if __name__ == '__main__':
    video_path = "datasets/video/0001_normal.mp4"
    image_folder = (video_path.split("/")[-1]).split(".")[0]
    image_path = "datasets/images/" + image_folder
    print(image_folder)
    try:
        shutil.rmtree(image_path)
        print("已经删除")
    except:
        print("已清空")
    finally:
        os.mkdir(image_path)

    video_to_frames(video_path=video_path, outPutDirName=image_path, img_pre=image_folder)
